#include <rclcpp/rclcpp.hpp>
#include <nav_msgs/msg/path.hpp>
#include <geometry_msgs/msg/pose_stamped.hpp>
#include <geometry_msgs/msg/pose.hpp>
#include <geometry_msgs/msg/transform_stamped.hpp>
#include <geometry_msgs/msg/twist.hpp>
#include <geometry_msgs/msg/pose_with_covariance_stamped.hpp>
#include <nav_msgs/msg/occupancy_grid.hpp>
#include <nav_msgs/msg/path.hpp>
#include <sensor_msgs/msg/point_cloud2.hpp>
#include <pcl/kdtree/kdtree_flann.h>
#include <pcl_conversions/pcl_conversions.h>
#include <pcl/filters/voxel_grid.h>
#include <pcl/point_types.h>
#include <pcl/point_cloud.h>
#include <pcl/common/transforms.h>
#include <pcl/filters/filter.h>
#include "tf2/LinearMath/Vector3.h"
#include <tf2_ros/static_transform_broadcaster.h>
#include <tf2_ros/buffer.h>
#include "tf2_ros/transform_listener.h"
#include "tf2_sensor_msgs/tf2_sensor_msgs.h"
#include <tf2_ros/create_timer_ros.h>
#include <tf2_ros/message_filter.h>
#include <tf2_geometry_msgs/tf2_geometry_msgs.hpp>
#include <cmath>

pcl::PointCloud<pcl::PointXYZ>::Ptr mapCloud(new pcl::PointCloud<pcl::PointXYZ>());
pcl::PointCloud<pcl::PointXYZ>::Ptr laserCloud(new pcl::PointCloud<pcl::PointXYZ>());
pcl::PointCloud<pcl::PointXYZ>::Ptr laserCloudCrop(new pcl::PointCloud<pcl::PointXYZ>());
pcl::PointCloud<pcl::PointXYZ>::Ptr laserCloudDwz(new pcl::PointCloud<pcl::PointXYZ>());
pcl::PointCloud<pcl::PointXYZI>::Ptr laserCloudI(new pcl::PointCloud<pcl::PointXYZI>());
//debug 
pcl::PointCloud<pcl::PointXYZ>::Ptr debugCloud(new pcl::PointCloud<pcl::PointXYZ>());

pcl::VoxelGrid<pcl::PointXYZ> laserDwzFilter;
pcl::KdTreeFLANN<pcl::PointXYZ> kdtree;
class LocalCostmap : public rclcpp::Node
{
  public:
    
    /* 雷达的坐标系 */
    std::string target_link="base_link";
    
    /*机器人半径，降采样用*/
    float robotRange=0.9;

    /*膨胀大小*/
    float inflatedRange=1.1;

    /*地图边长，一般设置为路面的宽度*2 */
    float costmapSize=9.0;

    /*最高点云*/
    float maxPointCloud=1.2;

    /*最低点云*/
    float minPointCloud=0.2;

    /*体素滤波*/
    double voxelGridResolution = 0.3;

    bool updateMap = false;

    // 地图缓存
    double origin_x, origin_y, resolution;
    int map_width, map_height;
    std::vector<int8_t> occupancy_grid;

    // 主类
    LocalCostmap() : Node("local_costmap")
    {
      publisher_local_costmap = this->create_publisher<nav_msgs::msg::OccupancyGrid>("local_costmap", 10);
      // 原始点云地图
      publisher_raw_costmap = this->create_publisher<nav_msgs::msg::OccupancyGrid>("raw_costmap", 10);
      map_sub = this->create_subscription<nav_msgs::msg::OccupancyGrid>(
              "/inflated_map", 
              10,
              std::bind(&LocalCostmap::topic_callback_map, 
              this, 
              std::placeholders::_1));

      rclcpp::QoS qos_profile(10);  // 创建一个QoS配置对象，队列大小为10

      // 设置历史记录策略为KEEP_LAST，并指定历史记录深度
      qos_profile.keep_last(10);

      cloud_sub = this->create_subscription<sensor_msgs::msg::PointCloud2>(
              "/segmentation/obstacle",
              10,
              std::bind(&LocalCostmap::topic_callback_cloud, 
              this, 
              std::placeholders::_1));

      tf_buffer_ = std::make_shared<tf2_ros::Buffer>(this->get_clock());
      tf_listener_ = std::make_shared<tf2_ros::TransformListener>(*tf_buffer_);
      // 初始化滤波器
      laserDwzFilter.setLeafSize(voxelGridResolution, voxelGridResolution, voxelGridResolution);  
      
    }

    void topic_callback_map(const nav_msgs::msg::OccupancyGrid & map)
    {    
      // 保证只更新一次局部地图
      if(!updateMap)
      {
        mapCloud->clear();
        /* map data */
        origin_x = map.info.origin.position.x;
        origin_y = map.info.origin.position.y;
        map_width = map.info.width;
        map_height = map.info.height;
        resolution = map.info.resolution;
        occupancy_grid=map.data;
        updateMap=true;
      }

    }

    void topic_callback_cloud(const sensor_msgs::msg::PointCloud2::ConstSharedPtr laserCloud2)
    {
      //等待地图接收
      if(!updateMap)
        return;

      // 地图消息初始化
      nav_msgs::msg::OccupancyGrid costmap;
      costmap.header.frame_id = "odom";
      costmap.info.resolution = resolution;
      costmap.info.width = costmapSize/resolution;
      costmap.info.height = costmapSize/resolution;
      costmap.data.resize(costmap.info.width * costmap.info.height, -1);
      
      // RCLCPP_INFO(this->get_logger(), "Node has been started.");
      // 清空点云
      laserCloud->clear();
      laserCloudDwz->clear();
      
      // 点云数据类型转换
      pcl::fromROSMsg(*laserCloud2, *laserCloud);

      // 点云预处理开始
      pcl::PointXYZ point;
      laserCloudCrop->clear();
      int laserCloudSize = laserCloud->points.size();

      // 范围降采样
      for (int i = 0; i < laserCloudSize; i++) {
        point = laserCloud->points[i];

        float pointX = point.x;
        float pointY = point.y;
        float pointZ = point.z;

        float dis = sqrt(pointX * pointX + pointY  * pointY);
        if ( dis > robotRange && dis<costmapSize && pointZ<maxPointCloud&&pointZ>minPointCloud) {
          point.x = pointX;
          point.y = pointY;
          point.z = pointZ;
          laserCloudCrop->push_back(point);
        }
      }

      // 坐标转换
      sensor_msgs::msg::PointCloud2 map_cloud_msg;
      pcl::toROSMsg(*laserCloudCrop, map_cloud_msg);
      try 
      {
          geometry_msgs::msg::TransformStamped transform = tf_buffer_->lookupTransform("odom", target_link, tf2::TimePointZero);

          costmap.info.origin.position.x = transform.transform.translation.x-costmapSize/2;
          costmap.info.origin.position.y = transform.transform.translation.y-costmapSize/2;

          tf2::doTransform(map_cloud_msg, map_cloud_msg, transform);
      } catch (tf2::TransformException& ex) {
          RCLCPP_ERROR(this->get_logger(), "Error transforming point cloud: %s", ex.what());
          return;
      }
      pcl::fromROSMsg(map_cloud_msg, *laserCloudCrop);
      
      // 体素过滤
      laserDwzFilter.setInputCloud(laserCloudCrop);
      laserDwzFilter.filter(*laserCloudDwz);

      // 遍历过滤后的点云，更新代价地图
      for (const auto& point : laserCloudDwz->points) 
      {
        // 防止出现无效点
        if (!pcl::isFinite<pcl::PointXYZ>(point)) 
        {
            continue;
        }

        // 将点云坐标转换为代价地图索引
        int map_x = static_cast<int>((point.x - costmap.info.origin.position.x) / resolution);
        int map_y = static_cast<int>((point.y - costmap.info.origin.position.y) / resolution);

        // 开始膨胀
        const int costmapWidth = costmap.info.width;
        const int costmapHeight = costmap.info.height;
        if (map_x >= 0 && map_x < costmapWidth - 1 && map_y >= 0 && map_y < costmapHeight - 1) {
            int index = map_y * costmapWidth + map_x;
            costmap.data[index] = 100; // 设置障碍物单元格的代价

            // 计算膨胀范围的起始和结束位置
            int x_start = std::max(0, map_x - static_cast<int>(inflatedRange / resolution));
            int x_end = std::min(costmapWidth, map_x + static_cast<int>(inflatedRange / resolution) + 1);

            for (int x = x_start; x < x_end; ++x) {
                // 计算垂直方向的起始和结束位置
                int j_start = std::max(0, map_y - static_cast<int>(std::sqrt(std::pow(inflatedRange / resolution, 2) - std::pow(x - map_x, 2))));
                int j_end = std::min(costmapHeight, map_y + static_cast<int>(std::sqrt(std::pow(inflatedRange / resolution, 2) - std::pow(x - map_x, 2))) + 1);

                for (int y = j_start; y < j_end; ++y) {
                    index = y * costmapWidth + x;
                    costmap.data[index] = 100; // 设置膨胀范围内的单元格代价
                }
            }
        }

      }
      publisher_raw_costmap->publish(costmap);

      // 将局部地图起始点坐标转换为代价地图索引开始值
      int global_map_x;
      int global_map_y;
      try 
      {
          geometry_msgs::msg::TransformStamped transform = tf_buffer_->lookupTransform("map", "odom", tf2::TimePointZero);
          global_map_x = static_cast<int>((costmap.info.origin.position.x - origin_x +transform.transform.translation.x) / resolution);
          global_map_y = static_cast<int>((costmap.info.origin.position.y - origin_y +transform.transform.translation.y) / resolution);

          tf2::doTransform(map_cloud_msg, map_cloud_msg, transform);
      } catch (tf2::TransformException& ex) {
          RCLCPP_ERROR(this->get_logger(), "Error transforming point cloud: %s", ex.what());
          return;
      }
      
      for(int x=0;x<costmapSize/resolution;x++)
      {
        for(int y=0;y<costmapSize/resolution;y++)
        {
          int global_index = (global_map_y + y) * map_width + (global_map_x + x);
          int local_index = y * costmap.info.width + x;
          if(global_index<map_width*map_height&&local_index < costmap.info.width * costmap.info.height)
          {
            if(occupancy_grid[global_index]==100)
              costmap.data[local_index] = 100;
          }
        }
      }

      publisher_local_costmap->publish(costmap);

      
    }

 
  private:
    //话题接收
    rclcpp::Subscription<nav_msgs::msg::OccupancyGrid>::SharedPtr map_sub;
    rclcpp::Subscription<sensor_msgs::msg::PointCloud2>::SharedPtr cloud_sub;
    rclcpp::Publisher<nav_msgs::msg::OccupancyGrid>::SharedPtr publisher_local_costmap;
    rclcpp::Publisher<nav_msgs::msg::OccupancyGrid>::SharedPtr publisher_raw_costmap;
    // Subscriber with tf2
    std::shared_ptr<tf2_ros::Buffer> tf_buffer_;
    std::shared_ptr<tf2_ros::TransformListener> tf_listener_;
    geometry_msgs::msg::TransformStamped transform;
};



int main(int argc, char *argv[])
{
    rclcpp::init(argc, argv);
    auto node = std::make_shared<LocalCostmap>();
    rclcpp::spin(node);
    rclcpp::shutdown();

    return 0;
}